from abc import ABC

from torch_geometric.data import Dataset as geometricDataset
from torch_geometric.data import Data
import pandas as pd
import os
import torch
from torch_geometric.data.data import BaseData

import const as cn


def get_load_dataset(data_path, task, data):
    def load_graph_dataset():
        if task == 'test':
            return ASTTestDataset(data_path, task, data)
        return ASTTrainDataset(data_path, task, data)

    return load_graph_dataset


class ASTDataset(geometricDataset, ABC):
    def get(self, idx: int) -> BaseData:
        return self.__getitem__(idx)

    def len(self) -> int:
        return self.data_len

    def __init__(self, data_path, task, data):
        super().__init__()
        self.data_path = data_path
        self.data = data
        self.data_list = data['file_id'].tolist()
        self.data_len = data.shape[0]
        self.task = task

    def __len__(self):
        return self.data_len


class ASTTrainDataset(ASTDataset):
    def __getitem__(self, index):
        act_index = self.data_list[index]
        _type = self.data.iloc[index].type

        pt_file_path = os.path.join(self.data_path, _type, str(act_index) + ".pt")
        dic = torch.load(pt_file_path)
        y = 0 if self.data.iloc[index].label == "white" else 1

        return Data(num_nodes=dic["node_type"].shape[0],
                    node_type=dic["node_type"],
                    node_value=dic["node_value"],
                    edge_index=dic["edge_index"].to(torch.int64),
                    edge_attr=dic["edge_attr"],
                    y=torch.tensor(y))


class ASTTestDataset(ASTDataset):
    def __getitem__(self, index):
        act_index = self.data_list[index]
        _type = self.data.iloc[index].type
        pt_file_path = os.path.join(self.data_path, _type, str(act_index) + ".pt")

        dic = torch.load(pt_file_path)
        return Data(num_nodes=dic["node_type"].shape[0],
                    node_type=dic["node_type"],
                    node_value=dic["node_value"],
                    edge_index=dic["edge_index"].to(torch.int64),
                    edge_attr=dic["edge_attr"])


